#!/bin/bash
# DeepGP 环境安装脚本 - 中国国内优化版
# 使用国内镜像源加速安装

echo "=========================================="
echo "  DeepGP 环境安装向导 (国内优化)"
echo "=========================================="
echo ""

# 获取项目根目录
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
PROJECT_ROOT="$( cd "$SCRIPT_DIR/.." && pwd )"
cd "$PROJECT_ROOT"

# 环境名称
ENV_NAME="deepgp"
PYTHON_VERSION="3.11"

# 检查 conda
if ! command -v conda &> /dev/null; then
    echo "❌ 未检测到 Conda"
    echo "请先安装 Anaconda 或 Miniconda"
    echo "下载地址: https://mirrors.tuna.tsinghua.edu.cn/anaconda/miniconda/"
    exit 1
fi

echo "✓ 检测到 Conda"

# 配置 conda 国内镜像源
echo ""
echo "配置 Conda 镜像源 (清华大学)..."
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/
conda config --set show_channel_urls yes
echo "✓ Conda 镜像源配置完成"

# 配置 pip 国内镜像源
echo ""
echo "配置 pip 镜像源 (清华大学)..."
mkdir -p ~/.pip
cat > ~/.pip/pip.conf << EOF
[global]
index-url = https://pypi.tuna.tsinghua.edu.cn/simple
[install]
trusted-host = pypi.tuna.tsinghua.edu.cn
EOF
echo "✓ pip 镜像源配置完成"

# 配置 HuggingFace 镜像
echo ""
echo "配置 HuggingFace 镜像..."
export HF_ENDPOINT=https://hf-mirror.com
echo "✓ HuggingFace 镜像: ${HF_ENDPOINT}"

# 检查环境是否已存在
if conda env list | grep -q "^${ENV_NAME} "; then
    echo ""
    echo "⚠️  环境 ${ENV_NAME} 已存在"
    read -p "是否删除并重新创建? (y/n) " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        echo "正在删除旧环境..."
        conda env remove -n ${ENV_NAME} -y
    else
        echo "使用现有环境"
    fi
fi

# 创建环境
if ! conda env list | grep -q "^${ENV_NAME} "; then
    echo ""
    echo "步骤 1/5: 创建 Conda 环境 (Python ${PYTHON_VERSION})"
    echo "=========================================="
    conda create -n ${ENV_NAME} python=${PYTHON_VERSION} -y
    if [ $? -ne 0 ]; then
        echo "❌ 创建环境失败"
        exit 1
    fi
    echo "✓ 环境创建成功"
fi

# 激活环境
echo ""
echo "步骤 2/5: 激活环境"
echo "=========================================="
source "$(conda info --base)/etc/profile.d/conda.sh"
conda activate ${ENV_NAME}

if [ $? -ne 0 ]; then
    echo "❌ 激活环境失败"
    exit 1
fi
echo "✓ 环境已激活: ${ENV_NAME}"

# 检测 GPU
echo ""
echo "步骤 3/5: 检测 GPU"
echo "=========================================="
if command -v nvidia-smi &> /dev/null; then
    echo "✓ 检测到 NVIDIA GPU"
    nvidia-smi --query-gpu=name --format=csv,noheader | head -n 1
    USE_GPU=true
else
    echo "⚠️  未检测到 GPU，将使用 CPU 版本"
    USE_GPU=false
fi

# 安装 PyTorch 2.1.2（DGL 兼容版本）
echo ""
echo "步骤 4/5: 安装 PyTorch 2.1.2 和 DGL"
echo "=========================================="
echo "使用 Conda 安装 PyTorch 2.1.2（国内最快）..."

if [ "$USE_GPU" = true ]; then
    # GPU 版本 - 优先 conda（国内最快）
    conda install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 pytorch-cuda=11.8 -c pytorch -c nvidia -y
    
    if [ $? -ne 0 ]; then
        echo "⚠️  Conda 安装失败，尝试清华 pip 源..."
        pip install torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 -i https://pypi.tuna.tsinghua.edu.cn/simple
    fi
else
    # CPU 版本
    conda install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 cpuonly -c pytorch -y || \
    pip install torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 -i https://pypi.tuna.tsinghua.edu.cn/simple
fi

echo "✓ PyTorch 2.1.2 安装完成"

# ⚠️ 重要：先安装兼容的 C++ 标准库（必须在 scipy 之前）
echo ""
echo "============================================"
echo "  安装兼容的 C++ 标准库"
echo "============================================"
conda install -c conda-forge libstdcxx-ng -y
echo "✓ C++ 标准库安装完成"

# 降级 NumPy 和 pandas（兼容性修复）
echo ""
echo "安装兼容版本的 NumPy 和 pandas..."
pip install "numpy<2.0" "pandas<2.0" -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装 DGL（与 PyTorch 2.1.2 兼容）
echo ""
echo "安装 DGL 1.1.3（与 PyTorch 2.1.2 兼容）..."

# 本地 DGL wheel 文件路径
DGL_WHEEL="$PROJECT_ROOT/whl/dgl-1.1.3+cu118-cp311-cp311-manylinux1_x86_64.whl"

echo ""
echo "============================================"
echo "  安装 DGL CPU 版本（无 CUDA 依赖）"
echo "============================================"

# 完全卸载可能存在的旧版本（避免 GPU 版本残留）
pip uninstall dgl -y 2>/dev/null
conda uninstall dgl -y 2>/dev/null
pip cache purge 2>/dev/null

# 删除可能残留的 DGL 文件
rm -rf /opt/miniconda3/envs/deepgp/lib/python3.11/site-packages/dgl* 2>/dev/null

# 方法1：使用 conda-forge（CPU 版本，最可靠）
echo "尝试 conda-forge 安装 DGL（推荐）..."
conda install -c conda-forge dgl -y

if [ $? -ne 0 ]; then
    echo ""
    echo "⚠️  conda-forge 失败，尝试官方 PyPI..."
    # 方法2：官方 PyPI（默认 CPU 版本）
    pip install dgl --no-cache-dir
fi

# 验证安装
echo ""
python -c "
import dgl
print(f'✓ DGL {dgl.__version__} 安装成功')
print(f'✓ DGL Backend: {dgl.backend.backend_name}')
" 2>/dev/null || {
    echo "⚠️  DGL 导入失败"
    echo "正在尝试最后的解决方案..."
    pip install dgl==1.1.3 --no-cache-dir
}

# 配置 DGL 后端
echo "配置 DGL 后端..."
python -c "import os; os.makedirs(os.path.expanduser('~/.dgl'), exist_ok=True); import json; json.dump({'backend': 'pytorch'}, open(os.path.expanduser('~/.dgl/config.json'), 'w')); print('✓ DGL 后端配置完成')" 2>/dev/null || echo "⚠️  DGL 配置失败（可忽略）"

# 安装其他依赖
echo ""
echo "步骤 5/5: 安装其他依赖"
echo "=========================================="

if [ -f "requirements.txt" ]; then
    echo "安装基础依赖..."
    pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
    echo "✓ 基础依赖安装完成"
fi

# ⚠️ 重要：重新安装 scipy 和 scikit-learn（确保使用新的 C++ 库）
echo ""
echo "重新安装 scipy 和 scikit-learn（确保兼容性）..."
pip install scipy scikit-learn -i https://pypi.tuna.tsinghua.edu.cn/simple --force-reinstall --no-deps
pip install scipy scikit-learn -i https://pypi.tuna.tsinghua.edu.cn/simple --no-deps

# ⚠️ 关键：最后再次确保 pandas 和 numpy 版本正确（防止被覆盖）
echo ""
echo "最终确保 pandas 和 numpy 版本正确..."
pip install "numpy<2.0" "pandas<2.0" -i https://pypi.tuna.tsinghua.edu.cn/simple --force-reinstall

if [ -f "streamlit_app/requirements.txt" ]; then
    echo "安装 Streamlit 依赖..."
    pip install -r streamlit_app/requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
    echo "✓ Streamlit 依赖安装完成"
fi

# 验证安装
echo ""
echo "=========================================="
echo "  验证安装"
echo "=========================================="

python << 'EOF'
import sys
import subprocess

# 使用 pip list 检查实际安装情况
def check_package_installed(pkg_name):
    try:
        result = subprocess.run(['pip', 'list'], capture_output=True, text=True)
        return pkg_name.lower() in result.stdout.lower()
    except:
        return False

packages = {
    'torch': 'PyTorch',
    'dgl': 'DGL',
    'transformers': 'Transformers',
    'pandas': 'Pandas',
    'numpy': 'NumPy',
    'fastNLP': 'fastNLP',
    'pyteomics': 'Pyteomics',
    'streamlit': 'Streamlit',
    'scikit-learn': 'scikit-learn',
    'wandb': 'Weights & Biases'
}

print("\n安装包检查:")
critical_missing = []
for package, name in packages.items():
    try:
        # 特殊处理某些包名
        import_name = package
        if package == 'scikit-learn':
            import_name = 'sklearn'
        
        mod = __import__(import_name)
        version = getattr(mod, '__version__', 'unknown')
        
        # 特别检查 pandas 版本（必须 < 2.0）
        if package == 'pandas':
            import pandas as pd
            major_version = int(version.split('.')[0]) if version != 'unknown' else 0
            if major_version >= 2:
                print(f"❌ {name:20s} {version} (版本过高！需要 < 2.0)")
                critical_missing.append(name)
                continue
        
        print(f"✓ {name:20s} {version}")
    except Exception as e:
        # 检查是否真的未安装
        if check_package_installed(package):
            print(f"⚠️  {name:20s} 已安装但导入失败: {str(e)[:50]}")
        else:
            print(f"❌ {name:20s} 未安装")
        
        if package in ['torch', 'dgl']:
            critical_missing.append(name)

# 检查 GPU 可用性
try:
    import torch
    print(f"\nPython 版本: {sys.version.split()[0]}")
    print(f"PyTorch 版本: {torch.__version__}")
    if torch.cuda.is_available():
        print(f"✓ GPU 可用: {torch.cuda.get_device_name(0)}")
        print(f"  GPU 数量: {torch.cuda.device_count()}")
        print(f"  CUDA 版本: {torch.version.cuda}")
    else:
        print("⚠️  使用 CPU 模式")
except Exception as e:
    print(f"⚠️  PyTorch 检查失败: {e}")

if critical_missing:
    print(f"\n❌ 关键依赖缺失或版本错误: {', '.join(critical_missing)}")
    if 'Pandas' in critical_missing:
        print("\n🔧 修复方案：")
        print("  pip install 'pandas<2.0' -i https://pypi.tuna.tsinghua.edu.cn/simple --force-reinstall")
    sys.exit(1)
else:
    print("\n✓✓✓ 所有依赖安装成功！")
EOF

if [ $? -eq 0 ]; then
    echo ""
    echo "=========================================="
    echo "  ✓✓✓ 安装完成！"
    echo "=========================================="
    echo ""
    echo "使用命令启动："
    echo "  ./run_app.sh"
    echo ""
else
    echo ""
    echo "=========================================="
    echo "  ❌ 安装失败"
    echo "=========================================="
    echo ""
    echo "请重新运行："
    echo "  conda remove -n deepgp --all -y"
    echo "  ./install.sh"
    echo ""
    exit 1
fi

